package com.liujit.gateway.gray;

import cn.hutool.core.util.StrUtil;
import com.liujit.common.constants.HeaderConstant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.loadbalancer.core.DelegatingServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.http.HttpHeaders;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @Description: 自定义服务实例筛选逻辑
 * @Author: liujun
 * @Date: 2021/7/1 7:33 下午
 */
@Slf4j
public class GrayServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier {

    public GrayServiceInstanceListSupplier(ServiceInstanceListSupplier delegate) {
        super(delegate);
    }

    @Override
    public Flux<List<ServiceInstance>> get() {
        return delegate.get();
    }

    @Override
    public Flux<List<ServiceInstance>> get(Request request) {
        return delegate.get(request).map(instances -> filteredByVersion(instances, getVersion(request.getContext())));
    }

    private List<ServiceInstance> filteredByVersion(List<ServiceInstance> instances, String version) {
        if (StrUtil.isEmpty(version)) {
            return instances;
        }
        List<ServiceInstance> filteredInstances = instances.stream()
                .filter(instance -> version.equalsIgnoreCase(instance.getMetadata().getOrDefault("version", "")))
                .collect(Collectors.toList());
        if (filteredInstances.size() > 0) {
            return filteredInstances;
        }
        return instances;
    }

    private String getVersion(Object requestContext) {
        if (requestContext == null) {
            return null;
        }
        if (requestContext instanceof RequestDataContext) {
            RequestDataContext context = (RequestDataContext) requestContext;
            if (context.getClientRequest() != null) {
                HttpHeaders headers = context.getClientRequest().getHeaders();
                if (headers != null) {
                    String version = headers.getFirst(HeaderConstant.VERSION);
                    log.info("request url is {}, request version is {}", context.getClientRequest().getUrl().getPath(),
                            version);
                    return version;
                }
            }
        }
        return null;
    }
}
